// Accelerated Dual Decomposition - An implementation of a MAP inference
// algorithm. See ai.stanford.edu/~vjojic/fastdd.pd
//
// Copyright 2011 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: nadavs@google.com (Nadav Samet)
//          noam.levy1@stanford.edu (Noam Levy)

#ifndef _ACCELERATED_DD_H_
#define _ACCELERATED_DD_H_

#include <stdio.h>
#include <vector>
#include <iostream>

using namespace std;

namespace AcceleratedDD {

  // Variable identifier (corresponds to i in the paper)
  typedef size_t  VariableId;

  // A scope is a vector of variable ids.
  typedef vector<VariableId> Scope;

  // Abstract class that Represents a factor
  class Factor {
   public:
    // Returns a list of variables in this factor's scope.
    virtual const Scope& GetScope() const = 0;

    virtual ~Factor() { };

    // Returns the number of distinct states the factor variables can take.
    // Denoted |Delta_c| in the paper.
    virtual size_t GetStateSize() const = 0;

    // Returns the number of binary variables needed to represent an
    // assignment.
    virtual size_t GetDimension() const = 0;

    // Evaluate the factor for the given assignment. Note that assignment
    // is a full assignment for all variables, not just to those that belong to
    // this scope.
    virtual double Evaluate(const size_t *assignment) const = 0;
  };

  // Represents an optimization problem, where the optimization objective
  // is the sum of the given factors.
  struct OptimizationProblem {
    size_t variables;        // number of variables
    size_t *variable_range;  // variable_range_[i] is |Val(X_i)|, the number of
                             // distinct labels that can be assigned to X_i.
    vector<Factor*> factors; // set of factors to optimize
    inline size_t factor_count() const { return factors.size(); }
  };

  // A concrete implementation of a factor that is backed up by an array. See
  // Factor for the description of all base class methods.
  class TableFactor : public Factor {
   public:
    // Constructs a new TableFactor with the given scope that is associated
    // with a given optimization problem.
    TableFactor(const Scope &scope,
                const OptimizationProblem *optimization_problem);

    virtual ~TableFactor();

    virtual const Scope& GetScope() const;

    virtual size_t GetStateSize() const;
    virtual size_t GetDimension() const;

    virtual double Evaluate(const size_t *assignment) const;

    // Assigns a values  to a single entry in this factor.
    // Params:
    //   scoped_assignment: array of |GetScope().size()| elements,
    //                      each corresponds to an assignment to the
    //                      corresponding variable in the scope.
    //   value:             value to assign.
    void Assign(const size_t *scoped_assignment, double value);

   private:
    const Scope scope_;
    const OptimizationProblem *optimization_problem_;
    size_t state_size_;

    // The entries are sorted in reverse lexicographical order.
    double *entries_;
  };

  // Three dimensional arrays of identical size.
  // In the paper terminology nu[c][i][a] is nu^c_{i,a}.
  // c is a factor index, i is a variable index in the scope of c, and a
  // represents a value that this variable can take.
  // For convenience, we allow i to take any variable index, even if it is
  // not a member of the factor phi_c.
  typedef double*** DualArray;

  // Structs that indicates that a given variable belongs to the factor
  // factor_index, and the variable index inside the factor is index.
  struct FactorPresence {
    size_t factor_index;
    size_t index;
  };

  typedef vector<FactorPresence> FactorPresenceList;

  struct AlgorithmState {
    // As in Algorithm 2 in the paper.
    // nu, eta and zeta are     double ***nu;
    DualArray nu;
    DualArray eta;
    DualArray zeta1, zeta2;
    DualArray *zeta, *zetaPrev;
    DualArray p;
    double theta;

    double *s_c;        // This is an array of length C (the number of factors)
                       // In the paper this is s^c_mu(nu_c).
    double dual_value;  // The sum of the above array, which is the current
                       // dual value.
    double nonsmooth_lower_bound;
    size_t dimension;        // Problem dimension
    int iteration;     // Current iteration

    // *presence_lists[var] is an array of FactorPresence for the variable var.
    // We initialize this structure in CreateInitialAlgorithmState.
    vector<FactorPresenceList > *presence_lists;
  };

  // Given an optimization problem, allocates and returns appropriately sized
  // initial state structure.
  AlgorithmState *CreateInitialAlgorithmState(const OptimizationProblem &op);

  // Deletes the given algorithm state.
  void DeleteAlgorithmState(const OptimizationProblem &p,
                            AlgorithmState *state);

  // Performs a single iteration of the algorithm starting from the given
  // state. The state variable gets updated.
  void Iterate(const OptimizationProblem &op, double mu, double L,
               AlgorithmState *state);

  // Given an assignment and a variable range vectors, increment the variables
  // in the scope in reverse lexicographic order (first variable in the scope
  // is least significant in the order)
  // Returns true when are back at (0, 0, ...).
  // This is used internally to iterate over all possible states of a scope.
  inline bool IncrementAssignment(size_t *assignment,
                                  const size_t *variable_range,
                                  const Scope &scope) {
    for (size_t i = 0; i < scope.size(); ++i) {
      if (assignment[scope[i]] < variable_range[scope[i]] - 1) {
        assignment[scope[i]]++;
        return true;
      }
      assignment[scope[i]] = 0;
    }
    return false;
  }

  enum LOGLEVEL {
    INNERMOST_LOOP = 0,     // noisy
    INLOOP = 1,
    INFO = 2,
    WARNING = 3,
    ERROR = 4,
    FATAL = 5
  };

  // Sets a stream where log messages will go (NULL, which is the default
  // means no logging
  void SetLogStream(ostream *logstream);

  // Sets the minimum loglevel (used to filter noisy messages from lower
  // levels)
  void SetLogLevel(LOGLEVEL loglevel);
}  // namespace AcceleratedDD

#endif  // _ACCELERATED_DD_H
